#include "iap.h"
#include "bsp/flash.h"
#include "bsp/board.h"
#include "iap.h"
#include "serial.h"
#include "utils/utils.h"
#include "protocol.h"

iap_t iap;

u32 iap_max_size(void)
{
	return CONFIG_APP_MAX_SIZE;
}

u32 iap_app_info_page(void)
{
	return CONFIG_PART_INFO;
}

void iap_write_begin(void)
{
	iap.address = PART_APP;
	iap.length = 0;
}

bool iap_write_check(u32 length, u32 checksum)
{
	if (!iap_flush(iap.length)) {
		return false;
	}

	if (iap_checksum(PART_APP, length) != checksum) {
		return false;
	}

	return true;
}

bool iap_write_end(u32 length, u32 checksum, u32 magic)
{
	iap_info_t info;

	if (!iap_write_check(length, checksum)) {
		return false;
	}

	info.length = length;
	info.checksum = checksum;
	info.magic = magic;

	return (bool) (flash_write_page(iap_app_info_page(), (u8 *)&info, sizeof(info), true) > 0);
}

bool iap_flush(u32 length)
{
	if (!flash_write_page(iap.address, iap.page, length, flash_iap_need_erase(iap.address))) {
		return false;
	}

	iap.address += length;
	iap.length = 0;

	return true;
}

bool iap_write(const u8 *buff, u16 length)
{
	while (1) {
		u16 remain = sizeof(iap.page) - iap.length;

		if (length < remain) {
			memcpy(iap.page + iap.length, buff, length);
			iap.length += length;
			break;
		}

		memcpy(iap.page + iap.length, buff, remain);
		if (!iap_flush(sizeof(iap.page))) {
			return false;
		}

		buff += remain;
		length -= remain;
	}

	return true;
}

bool iap_write_u32(u32 value)
{
	return iap_write((u8 *) &value, sizeof(value));
}

#if CONFIG_IAP_CRC32
u32 iap_checksum_init(void)
{
	return 0xFFFFFFFFU;
}

u32 iap_checksum_finish(u32 crc)
{
	return crc ^ 0xFFFFFFFFU;
}

u32 iap_checksum_update(u32 crc, const u8 *buff, u32 length)
{
	const u8 *buff_end;

	for (buff_end = buff + length; buff < buff_end; buff++) {
		u8 value = *buff;
		u8 i;

		for (i = 0; i < 8; i++) {
			if (((crc ^ value) & 1) != 0) {
				crc = (crc >> 1) ^ 0xEDB88320U;
			} else {
				crc >>= 1;
			}

			value >>= 1;
		}
	}

	return crc;
}
#else
u32 iap_checksum_init(void)
{
	return 0;
}

u32 iap_checksum_finish(u32 checksum)
{
	return checksum;
}

u32 iap_checksum_update(u32 checksum, const u8 *buff, u32 length)
{
	const u8 *buff_end = buff + length;

	for (buff_end = buff + length; buff < buff_end; buff++) {
		checksum += *buff;
	}

	return checksum;
}
#endif

u32 iap_checksum(u32 addr, u32 length)
{
	u32 checksum = iap_checksum_init();
	checksum = iap_checksum_update(checksum, (u8 *) addr, length);
	return iap_checksum_finish(checksum);
}

#ifdef CONFIG_FMC_UPGRADE
bool iap_fmc_upgrade(void)
{
	const iap_info_t *info = (iap_info_t *) iap_last_page();
	u32 position;

	if (info->success != MAGIC_FMC_UPGRADE) {
		return false;
	}

	if (info->address == 0xFFFFFFFF) {
		return true;
	}

	if (info->length > iap_max_size()) {
		return true;
	}

	if (iap_checksum(info->address, info->length) != info->checksum) {
		return true;
	}

	position = 0;

	iap_write_begin();

	while (position < info->length) {
		u32 remain = info->length - position;
		u8 buff[256];

		if (remain > sizeof(buff)) {
			remain = sizeof(buff);
		}

		memcpy(buff, (void *) (info->address + position), remain);
		iap_write(buff, remain);
		position += remain;

		wdg_reload();
	}

	return iap_write_check(position, info->checksum);
}
#endif

#ifndef CONFIG_BOOT_UPDATE

// ================================================================================

__STATIC_INLINE int iap_process_cmd_write_begin(const u8 *buff, u16 length)
{
	iap.write_position = 0;
	iap_write_begin();
	return 0;
}

__STATIC_INLINE int iap_process_cmd_write_end(const u8 *buff, u16 length)
{
	u32 size, checksum;

#ifndef CONFIG_REDUCE_STORAGE
	if (length < 7) {
		return -1;
	}
#endif

	size = decode_u24((u8 *)buff);
	checksum = decode_u32((u8 *)buff + 3);

	if (iap_write_end(size, checksum, MAGIC_APP_READY)) {
		return 0;
	}

	return -1;
}

__STATIC_INLINE int iap_process_cmd_write_slow(u8 *buff, u16 length)
{
	u32 position;

#ifndef CONFIG_REDUCE_STORAGE
	if (length < 3) {
		return -1;
	}
#endif

	position = decode_u24(buff);

	if (iap.write_position == position) {
		length -= 3;

		if (iap_write(buff + 3, length)) {
			iap.write_position += length;
		} else {
			return -1;
		}
	}

	encode_u24(buff + 1, iap.write_position);

	return 3;
}

#ifndef CONFIG_REDUCE_STORAGE
__STATIC_INLINE int iap_process_cmd_read_data(u8 *args, u16 length)
{
	u32 addr;

	if (length < 4) {
		return -1;
	}

	addr = decode_u24(args);
	length = args[3];

	if (length > 250) {
		length = 250;
	}

	encode_u24(args + 1, addr);
	memcpy(args + 4, (void *) (CONFIG_FLASH_ADDR + addr), length);

	return length + 3;
}
#endif

__STATIC_INLINE int iap_process_cmd_read_string(u8 *args, u16 length)
{
	const char *text;
	u32 addr;

#ifndef CONFIG_REDUCE_STORAGE
	if (length < 3) {
		return -1;
	}
#endif

	addr = decode_u24(args);
	u32 offset = addr;
#if defined(STM32F4)
	/* 兼容16K的bootloader + 16K info*/
	if (addr == 0x2A00) {
		offset = 0x8A00;
	}else if (addr == 0x2800) {
		offset = 0x8800;
	}else if (addr == 0x2C00) {
		offset = 0x8C00;
	}
#elif defined (STM32G4)
#if defined(CONFIG_BOARD_USB2CAN)
	/* 兼容16K的bootloader + 4K info*/
	if (addr == 0x2A00) {
		offset = 0x5A00;
	}else if (addr == 0x2800) {
		offset = 0x5800;
	}else if (addr == 0x2C00) {
		offset = 0x5C00;
	}
#else
	/* 兼容8K的bootloader + 4K info*/
	if (addr == 0x2A00) {
		offset = 0x3A00;
	}else if (addr == 0x2800) {
		offset = 0x3800;
	}else if (addr == 0x2C00) {
		offset = 0x3C00;
	}
#endif /* CONFIG_BOARD_USB2CAN */
#endif /* STM32F4 */

	text = (const char *) (CONFIG_FLASH_ADDR + offset);

	for (length = 0; (text[length] != (char)0xFF) && text[length] && (length < 42); length++);

	encode_u24(args + 1, addr);
	memcpy(args + 4, text, length);

	return length + 3;
}

__STATIC_INLINE void iap_process_command(can_pack_t *pack)
{
	int rsplen;
	u16 length;
	u8 *args;
	u16 key;

#ifndef CONFIG_REDUCE_STORAGE
	if (length < 2) {
		return;
	}
#endif
	key = *(u16 *) pack->data;
	length = pack->length - 2;
	args = pack->data + 2;

	switch (key)
	{
	case CAN_BUILD_CMD(0xF0):
		iap_cancel_runapp();
		iap.state = IAP_ENTER;
		rsplen = 0;
		break;

	case CAN_BUILD_CMD(0xF1):
		iap.state = IAP_BEGIN;
		rsplen = iap_process_cmd_write_begin(args, length);
		break;

	case CAN_BUILD_CMD(0xF2):
		iap.state = IAP_CHECK;
		rsplen = iap_process_cmd_write_end(args, length);
		break;

#ifndef CONFIG_REDUCE_STORAGE
	case CAN_BUILD_CMD(0xF3):
		iap.state = IAP_WRITE;
		iap_write(args, length);
		return;
#endif

	case CAN_BUILD_CMD(0xF4):
		iap.state = IAP_WRITE;
		rsplen = iap_process_cmd_write_slow(args, length);
		break;

	case CAN_BUILD_CMD(0xF5):
		iap.state = IAP_BOOT;
		if (iap_boot(false)) {
			rsplen = 0;
		} else {
			rsplen = -1;
		}
		break;

	case CAN_BUILD_CMD(0xF6):
		args[1] = iap.state;
		rsplen = 1;
		break;

#ifndef CONFIG_REDUCE_STORAGE
	case CAN_BUILD_CMD(0xF7):
		rsplen = iap_process_cmd_read_data(args, length);
		break;
#endif

	case CAN_BUILD_CMD(0xF8):
		rsplen = iap_process_cmd_read_string(args, length);
		break;

	default:
		rsplen = -1;
		break;
	}

	if (rsplen < 0) {
		rsplen = 0;
		args[0] = 0x01;
	} else {
		args[0] = 0x00;
		iap.ticks = 0;
	}
#if (CONFIG_SERIAL_TYPE == SERIAL_BOOTLOADER)
	serial_send_response(pack->src, pack->dest, pack->data, rsplen + 3);
#else
	can_send_response(pack->src, pack->dest == 0x7F ? CONFIG_ADDR_ID : pack->dest, pack->data, rsplen + 3);
#endif
}

void iap_process_pack(can_pack_t *pack)
{
	iap_process_command(pack);
}

void iap_run_app_delay(u32 ticks) {
	wdog_init();
	iap.run_app_delay = ticks;
}

void iap_cancel_runapp(void) {
	iap.run_app_delay = -1;
}

#define stack_pointer() (REG32(PART_APP))
#define app_entry()     (REG32(PART_APP + 4))
static bool can_try_booting(void)
{
    if (app_entry() != 0xFFFFFFFF &&  stack_pointer()!= 0xFFFFFFFF) {
		return true;
	}
	return false;
}

void iap_running_app(void) {
	SysTick->CTRL = 0;

	bsp_deinit();

	/* initialize user application's stack pointer */
	__set_MSP(stack_pointer());
	nvic_vector_table_set(PART_APP, 0);

	/* jump to user application */
	((void (*)(void))app_entry())();
}

bool iap_boot(bool normal)
{
	const iap_info_t *info = (iap_info_t *) iap_app_info_page();
	const char *board_name = (const char *) PART_NAME;

	iap.run_app_delay = -1;

	if (strcmp(board_name, CONFIG_BOARD_NAME) != 0) {
		return false;
	}
	if (info->magic == MAGIC_APP_BOOT_OK) {
		iap_running_app();
		return false;
	}
	if (info->magic == MAGIC_DUAL_BANK_UPGRADE) {
		dual_iap_info_t *dual_info = (dual_iap_info_t *)iap_app_info_page();
		if (dual_info->dual_checksum == iap_checksum(dual_info->dual_addr, dual_info->dual_len)) {
			int buf_size = sizeof(iap.page);
			int remain_len = dual_info->dual_len;
			u8 *buffer = &iap.page[0];
			u32 src = dual_info->dual_addr;
			u32 dst = CONFIG_PART_APP;
			while(remain_len > 0) {
				int op_len = remain_len>buf_size?buf_size:remain_len;
				memcpy(buffer, (void *)src, op_len);
				flash_write_page(dst, buffer, op_len, flash_iap_need_erase(dst));
				dst += op_len;
				src += op_len;
				remain_len -= op_len;
			}
			if (dual_info->dual_checksum == iap_checksum(CONFIG_PART_APP, dual_info->dual_len)) {
				dual_iap_info_t w_info;
				w_info.info.checksum = dual_info->dual_checksum;
				w_info.info.length = dual_info->dual_len;
				w_info.info.magic = MAGIC_APP_READY;
				w_info.info.address = CONFIG_PART_APP;
				flash_write_page(CONFIG_PART_INFO, (u8 *)&w_info, sizeof(w_info), true);
				iap_run_app_delay(CONFIG_TRY_BOOT_DELAY);
			}
		}else if (info->length < CONFIG_APP_MAX_SIZE) {
			if (info->checksum == iap_checksum(PART_APP, info->length)){
				if (can_try_booting()) {
					iap_run_app_delay(CONFIG_TRY_BOOT_DELAY);
				}
			}
		}
	}else if (info->magic != MAGIC_ENTER_BOOT) {
		if (info->magic == 0xFFFFFFFF && info->address == 0xFFFFFFFF && info->checksum == 0xFFFFFFFF && info->length == 0xFFFFFFFF) {
			if (can_try_booting()) {
				iap_run_app_delay(CONFIG_TRY_BOOT_DELAY);
			}
		}else if (info->length < CONFIG_APP_MAX_SIZE) {
			if (info->checksum == iap_checksum(PART_APP, info->length)){
				if (can_try_booting()) {
					iap_run_app_delay(CONFIG_TRY_BOOT_DELAY);
					if(!normal) {
						return true;
					}
				}
			}
		}
	}
	return false;
}
#endif

// ================================================================================
#if 0
static bool iap_write_bootloader(void)
{
	const u8 *buff = (const u8 *) PART_UPDATE;
	u32 addr = PART_BOOT;
#ifndef CONFIG_BATT_UPDATE
	iap_info_t info;
#endif
	u32 checksum;
	u32 size;

	size = decode_u24(buff - 3);
	if (size > 0x2000) {
		return false;
	}

	checksum = decode_u32(buff - 7);
	if (checksum != iap_checksum(PART_UPDATE, size)) {
		return false;
	}

	if (strcmp("Boot", (const char *) (buff + 0x0C00)) != 0) {
		return false;
	}

#ifndef CONFIG_BATT_UPDATE
	memcpy(&info, (iap_info_t *) PART_FACTORY, sizeof(info));
#endif

	while (size > 0) {
		u16 length;

		if (size > FMC_PAGE_SIZE) {
			length = FMC_PAGE_SIZE;
		} else {
			length = size;
		}

		memcpy(iap.page, buff, length);
		fmc_write_page(addr, iap.page, length, true);

		buff += length;
		addr += length;
		size -= length;
	}

#ifndef CONFIG_BATT_UPDATE
	fmc_write_page(PART_FACTORY, (u8 *) &info, sizeof(info), false);
#endif

	return true;
}

void iap_upgrade_bootloader(void)
{
	iap_write_bootloader();

	fmc_unlock();
	fmc_page_erase(iap_last_page());
	fmc_lock();

	NVIC_SystemReset();
}
#endif
